package com.webchat.search.service.voctor;


import com.webchat.common.bean.APIResponseBean;
import com.webchat.common.bean.APIResponseBeanUtil;
import com.webchat.common.exception.AiCallException;
import com.webchat.common.util.JsonUtil;
import com.webchat.rmi.aigc.BasicModelServiceClient;
import io.milvus.v2.client.MilvusClientV2;
import io.milvus.v2.common.IndexParam;
import io.milvus.v2.service.collection.request.CreateCollectionReq;
import io.milvus.v2.service.collection.request.DropCollectionReq;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;

import java.util.List;


@Slf4j
public abstract class AbstractMilvusService<T, R> {

    @Autowired
    public MilvusClientV2 client;

    @Autowired
    private BasicModelServiceClient basicModelClient;

    /**
     * 默认查询相似度最高的1条
     */
    private final static int DEFAULT_TOP_K = 1;

    private final static int DEFAULT_SCORE = 0;

    /**
     * 是否需要走队列模式，默认不需要
     *
     * @return
     */
    protected boolean queueModel() {
        return false;
    }

    /**
     * 加入队列，普通业务场景不需要单独实现，队列消费场景自己实现
     * @param data
     * @return
     */
    protected boolean addTaskQueue(T data) {
        return true;
    }

    /**
     * 向量纬度
     * @return
     */
    protected abstract int dimension();

    /**
     * 集合名称，类似RDBMS中的数据库概念
     *
     * 明明定义由具体实现类来定义
     * @return
     */
    protected abstract String collectionName();

    /**
     * schema 定义
     *
     * @return
     */
    protected abstract CreateCollectionReq.CollectionSchema schema();

    /**
     * 集合schema索引定义
     *
     * @return
     */
    protected abstract List<IndexParam> indexParams();

    /**
     * 数据写入
     *
     * @param data
     * @return
     */
    public abstract boolean writeCollection(T data);

    /**
     * 处理同步任务数据
     * @param data
     */
    public void doHandleSyncData(T data) {

        boolean queueMedal = this.queueModel();
        if (queueMedal) {
            // 队列模式处理
            this.addTaskQueue(data);
        } else {
            this.writeCollection(data);
        }
    }

    /**
     * 基于相似度的 TopK 条数据搜索
     * @param query
     * @param topK
     * @return
     */
    public abstract List<R> search(String query, int topK, float gtScore);

    public List<R>  search(String query, int topK) {
        return this.search(query, topK, DEFAULT_SCORE);
    }

    public List<R>  search(String query) {

        return this.search(query, DEFAULT_TOP_K);
    }

    /**
     * 单条文本 embedding（文本转向量）
     *
     * @param text
     * @return
     */
    protected float[] embed(String text) {

        APIResponseBean<float[]> response = basicModelClient.embed(text);
        if (APIResponseBeanUtil.isOk(response)) {
            return response.getData();
        }
        log.error("embeeding err ===> text:{}, response:{}", text, JsonUtil.toJsonString(response));
        throw new AiCallException("embedding err!");
    }

    /**
     * 批量文本 embedding（文本转向量）
     *
     * @param texts
     * @return
     */
    protected List<float[]> embed(List<String> texts) {

        APIResponseBean<List<float[]>> response = basicModelClient.embed(texts);
        if (APIResponseBeanUtil.isOk(response)) {
            return response.getData();
        }
        log.error("embeeding err ===> texts:{}, response:{}",
                JsonUtil.toJsonString(texts), JsonUtil.toJsonString(response));
        throw new AiCallException("embedding err!");
    }

    /**
     * 创建集合
     */
    public void createCollection() {
        /**
         * 集合名
         */
        String collectionName = this.collectionName();
        /**
         * 定义schema
         */
        CreateCollectionReq.CollectionSchema schema = this.schema();
        /**
         * 索引
         */
        List<IndexParam> indexParams = this.indexParams();
        /**
         * 创建集合
         */
        CreateCollectionReq request = CreateCollectionReq.builder()
                .collectionName(collectionName)
                .collectionSchema(schema)
                .indexParams(indexParams)
                .build();
        client.createCollection(request);
    }

    /**
     * 删除集合
     *
     */
    public void dropCollection() {
        DropCollectionReq request = DropCollectionReq.builder()
                .collectionName(this.collectionName())
                .build();
        client.dropCollection(request);
    }
}
